// @HEADER
//*********************************************************************//
//  SiYuan: A numerical PDE solver                                     //
//  Copyright (2022) YUAN Xi                                           //
//  This Software is released under the BSD 2-Clause license detailed  //
//  in the file "LICENSE" in the top-level SiYuan directory            //
//*********************************************************************//
// @HEADER

#ifndef _BC_DIRICHLET_IMPL_HPP
#define _BC_DIRICHLET_IMPL_HPP

#include "Panzer_Workset.hpp"
#include "Panzer_GlobalEvaluationDataContainer.hpp"
#include "Phalanx_DataLayout_MDALayout.hpp"

namespace SiYuan {

//**********************************************************************
template<typename EvalT>
Dirichlet<EvalT>::Dirichlet(const Teuchos::ParameterList& p, 
 const Teuchos::RCP<const panzer_stk::STK_Interface>& mesh, 
 const Teuchos::RCP<const panzer::GlobalIndexer> & indexer )
: SomeCondition()
{
    Teuchos::RCP<Teuchos::ParameterList> params = Teuchos::parameterList();
    *params = p;

  // ********************
  // Validate and parse parameter list
  // ********************
  {    
    Teuchos::ParameterList valid_params;
    valid_params.set<std::string>("Method", "1_0");
    valid_params.set<std::string>("Nodeset Name", "");
    valid_params.set<std::string>("Sideset Name", "");
    valid_params.set<std::string>("Type", "");
    valid_params.set<int>("Group ID", 0);
    valid_params.set<double>("Penalty", 1.e30);
    valid_params.set< Teuchos::Array<std::string> >("Variable Names", Teuchos::tuple<std::string>("") );
    valid_params.sublist("Data").disableRecursiveValidation();

    params->validateParametersAndSetDefaults(valid_params);
  }
    m_method = params->get<std::string>("Method");
    m_nodeset_name = params->get<std::string>("Nodeset Name");
    m_edgeset_name = params->get<std::string>("Sideset Name");
    m_type = params->get<std::string>("Type");
    m_penalty = params->get<double>("Penalty");
    m_group_id = params->get<int>("Group ID");
    m_dof_name = params->get< Teuchos::Array<std::string> >("Variable Names");
//indexer->print_DOFInfo(std::cout);
    auto pl_val = params->sublist("Data");
    m_value = XYZLib::variable<double>(pl_val);

    m_local_dofs.clear();
    if( !m_nodeset_name.empty() ) 
    {
      std::vector<stk::mesh::EntityId> nodes;
      mesh->getAllNodeSetIds(m_nodeset_name,nodes);
      for(auto cnd: nodes) {
        m_global_nodes.insert(cnd);
      }
      
      for(auto myname: m_dof_name) {
        int fdnum = indexer->getFieldNum(myname);
        for ( auto nd: m_global_nodes ) {
          auto b = indexer->getNodalLDofOfField( fdnum, nd );
          if( b<0 ) std::cout << fdnum << ", " << nd <<std::endl;
          TEUCHOS_TEST_FOR_EXCEPTION( (b<0), std::logic_error,
				    "Error - Cannot find dof of Nodeset!" );
          m_local_dofs.emplace_back(b);
        }
        // globalIndexer_ -> getNodesetsLocalIndex( fdnum, nnodes, dofs);
      }
    }

    if( !m_edgeset_name.empty() ) 
    {
    //  std::vector<stk::mesh::Entity>  sides;
    //  mesh->getMySides(m_edgeset_name,sides); 

      std::vector<stk::mesh::EntityId> edges;
    //  mesh->getAllEdgesId(m_edgeset_name,edges); 

      mesh->getAllEdgeSetIds(m_edgeset_name,edges); 

    //  m_local_dofs.clear();
      for(auto myname: m_dof_name) {
        int fdnum = indexer->getFieldNum(myname);
        for ( auto nd: edges ) {
          auto b = indexer->getEdgeLDofOfField( fdnum, nd );
          if( b.empty() ) std::cout << fdnum << ", " << nd <<std::endl;
          TEUCHOS_TEST_FOR_EXCEPTION( (b.empty()), std::logic_error,
				    "Error - Cannot find dof of Edgeset!" );
      //    m_local_dofs.emplace_back(b);
        }
      }
    }
    
};

//=======================================================================
template< typename ScalarT >
void 
readDirichletCondition( std::vector< Dirichlet<ScalarT> >& bcs, const Teuchos::ParameterList& pl,
    const Teuchos::RCP<const panzer_stk::STK_Interface>& mesh,
    const Teuchos::RCP<const panzer::GlobalIndexer> & indexer )
{
  bcs.clear();
  
  std::size_t bc_index = 0;
  for (Teuchos::ParameterList::ConstIterator bc_pl=pl.begin(); bc_pl != pl.end(); ++bc_pl,++bc_index) {
    TEUCHOS_TEST_FOR_EXCEPTION( !(bc_pl->second.isList()), std::logic_error,
				"Error - All objects in the Dirichlet Conditions sublist must be sublists!" );
    Teuchos::ParameterList& sublist = Teuchos::getValue<Teuchos::ParameterList>(bc_pl->second);
    
    Dirichlet<ScalarT> bc(sublist,mesh,indexer);
    bcs.emplace_back(bc);
  }
}

//**********************************************************************
template<typename Traits>
DirichletsEvalutor<panzer::Traits::Residual, Traits> :: DirichletsEvalutor( const std::vector< Dirichlet<ScalarT> >& ds,
  const Teuchos::RCP<const panzer_stk::STK_Interface>& mesh)
: PHX::EvaluatorWithBaseImpl<Traits>("Dirichlet Boundary Conditions")
{
  //std::vector<stk::mesh::EntityId> nodes;

  m_dirichlets.clear();
  for( auto d: ds )
  {
    if( !d.isActive() ) continue;
    if( d.getType()=="tokamak") {
	    const auto & vals = d.getValues();
    //  const auto & name = d.getNodeSetName();
    //  mesh->getAllNodeSetIds(name,nodes); 
      std::size_t cc=0;
      for( auto nd: d.m_global_nodes)
      {
        const double * coord = mesh->getNodeCoordinates(nd);
        double r=coord[0];
        double z=coord[1];
        double v =  (1. -4.* z* z -(2.*( r -1.) +(r -1.)*(r -1.)) *(2.*( r -1.) +(r -1.)*(r -1.))) /8.;
      //  double v=0.5*r0*r0*a*a*(1.0-z*z/(a*a)-((r-r0)/a+(r-r0)*(r-r0)/(2.0*a*r0))*((r-r0)/a+(r-r0)*(r-r0)/(2.0*a*r0)));
      //  double v=0.25*r -0.125*r*r-0.5*z;std::cout << v << "," << r << "," << z << std::endl;
      //  double v = (-0.125*vals[0]+vals[3])*r -0.5*vals[1]*z + vals[2]+vals[4]*(r*r-4.0*r*z);
        const auto& dof = d.m_local_dofs[cc]; cc++;
        m_dirichlets.insert( std::make_pair(dof, v) );
      }
    } else {
      auto v = d.getValue();
      for( auto dof: d.m_local_dofs )
      {
        m_dirichlets.insert( std::make_pair(dof, v) );
      }
    }
  }

  Teuchos::RCP<PHX::DataLayout> dummy = Teuchos::rcp(new PHX::MDALayout<void>(0));
  const PHX::Tag<ScalarT> fieldTag("Dirichlet Condition", dummy);
  this->addEvaluatedField(fieldTag);
  this->setName("Dirichlet_Residual");
};

template<typename Traits>
void DirichletsEvalutor<panzer::Traits::Residual, Traits> :: preEvaluate(typename Traits::PreEvalData d)
{
 // typedef panzer::TpetraLinearObjContainer<ScalarT,panzer::LocalOrdinal,panzer::GlobalOrdinal> LOC;

  // extract linear object container
  // m_GhostedContainer = Teuchos::rcp_dynamic_cast<LOC>(d.gedc->getDataObject("Ghosted Container"));

  if(Teuchos::is_null(m_GhostedContainer))
    m_GhostedContainer = Teuchos::rcp_dynamic_cast<panzer::LinearObjContainer>(d.gedc->getDataObject("Ghosted Container"));
}

template<typename Traits>
void DirichletsEvalutor<panzer::Traits::Residual, Traits> :: evaluateFields(typename Traits::EvalData d)
{
/*  std::vector<panzer::LocalOrdinal> indx;
  for (auto i = m_dirichlets.begin(); i != m_dirichlets.end(); ++i){
    indx.emplace_back(i->first);
  }*/
  m_GhostedContainer->evalDirichletResidual(m_dirichlets);
}

//**********************************************************************
template<typename Traits>
DirichletsEvalutor<panzer::Traits::Jacobian, Traits> :: DirichletsEvalutor( const std::vector< Dirichlet<ScalarT> >& ds,
  const Teuchos::RCP<const panzer_stk::STK_Interface>& mesh)
: PHX::EvaluatorWithBaseImpl<Traits>("Dirichlet Boundary Conditions")
{
  //std::vector<stk::mesh::EntityId> nodes;

  m_dirichlets.clear();
  for( auto d: ds )
  {
    if( !d.isActive() ) continue;
    if( d.getType()=="tokamak") {
	    const auto & vals = d.getValues();
    //  const auto & name = d.getNodeSetName();
    //  mesh->getAllNodeSetIds(name,nodes); 
      std::size_t cc=0;
      for( auto nd: d.m_global_nodes)
      {
        const double * coord = mesh->getNodeCoordinates(nd);
        double r=coord[0];
        double z=coord[1];
        double v =  (1. -4.* z* z -(2.*( r -1.) +(r -1.)*(r -1.)) *(2.*( r -1.) +(r -1.)*(r -1.))) /8.;
      //  if( z==-0.8) std::cout << "b:" << r << ", " << v << std::endl;
      //  double v = (-0.125*vals[0]+vals[3])*r -0.5*vals[1]*z + vals[2]+vals[4]*(r*r-4.0*r*z);
        const auto& dof = d.m_local_dofs[cc]; cc++;
        m_dirichlets.insert( std::make_pair(dof, v) );
      }
    } else {
      auto v = d.getValue();
      for( auto dof: d.m_local_dofs )
      {
        m_dirichlets.insert( std::make_pair(dof, v) );
      }
    }
  }

  Teuchos::RCP<PHX::DataLayout> dummy = Teuchos::rcp(new PHX::MDALayout<void>(0));
  const PHX::Tag<ScalarT> fieldTag("Dirichlet Condition", dummy);
  this->addEvaluatedField(fieldTag);
  this->setName("Dirichlet_Jacobian");
};

template<typename Traits>
void DirichletsEvalutor<panzer::Traits::Jacobian, Traits> :: preEvaluate(typename Traits::PreEvalData d)
{
  m_GhostedContainer = Teuchos::rcp_dynamic_cast<panzer::LinearObjContainer>(d.gedc->getDataObject("Ghosted Container"));
}

template<typename Traits>
void DirichletsEvalutor<panzer::Traits::Jacobian, Traits> :: evaluateFields(typename Traits::EvalData d)
{
//  m_GhostedContainer->writeMatrixMarket("out.mm");
//  if( m_method == DirichletStrategy::I1_0 ) 
//    m_GhostedContainer->applyDirichletBoundaryCondition(m_dirichlets);
//  else  //penalty
//  {
    double val = d.pivot_dirichlet;
    m_GhostedContainer->applyDirichletBoundaryCondition(val, m_dirichlets);
//  }
  
}

//**********************************************************************
template<typename Traits>
DirichletsEvalutor<panzer::Traits::Tangent, Traits> :: DirichletsEvalutor( const std::vector< Dirichlet<ScalarT> >& ds,
  const Teuchos::RCP<const panzer_stk::STK_Interface>& mesh)
: PHX::EvaluatorWithBaseImpl<Traits>("Dirichlet Boundary Conditions")
{
  //std::vector<stk::mesh::EntityId> nodes;

  m_dirichlets.clear();
  for( auto d: ds )
  {
    if( !d.isActive() ) continue;
    if( d.getType()=="tokamak") {
	    const auto & vals = d.getValues();
    //  const auto & name = d.getNodeSetName();
    //  mesh->getAllNodeSetIds(name,nodes); 
      std::size_t cc=0;
      for( auto nd: d.m_global_nodes)
      {
        const double * coord = mesh->getNodeCoordinates(nd);
        double r=coord[0];
        double z=coord[1];
        double v =  (1. -4.* z* z -(2.*( r -1.) +(r -1.)*(r -1.)) *(2.*( r -1.) +(r -1.)*(r -1.))) /8.;
      //  if( z==-0.8) std::cout << "a:" << r << ", " << v << std::endl;
      //  double v = (-0.125*vals[0]+vals[3])*r -0.5*vals[1]*z + vals[2]+vals[4]*(r*r-4.0*r*z);
        const auto& dof = d.m_local_dofs[cc]; cc++;
        m_dirichlets.insert( std::make_pair(dof, v) );
      }
    } else {
      auto v = d.getValue();
      for( auto dof: d.m_local_dofs )
      {
        m_dirichlets.insert( std::make_pair(dof, v) );
      }
    }
  }

  Teuchos::RCP<PHX::DataLayout> dummy = Teuchos::rcp(new PHX::MDALayout<void>(0));
  const PHX::Tag<ScalarT> fieldTag("Dirichlet Condition", dummy);
  this->addEvaluatedField(fieldTag);
  this->setName("Dirichlet_Tangent");
};

template<typename Traits>
void DirichletsEvalutor<panzer::Traits::Tangent, Traits> :: preEvaluate(typename Traits::PreEvalData d)
{
  m_GhostedContainer = Teuchos::rcp_dynamic_cast<panzer::LinearObjContainer>(d.gedc->getDataObject("Ghosted Container"));
}

template<typename Traits>
void DirichletsEvalutor<panzer::Traits::Tangent, Traits> :: evaluateFields(typename Traits::EvalData d)
{}

#ifdef Panzer_BUILD_HESSIAN_SUPPORT
//**********************************************************************
template<typename Traits>
DirichletsEvalutor<panzer::Traits::Hessian, Traits> :: DirichletsEvalutor( const std::vector< Dirichlet<ScalarT> >& ds)
: PHX::EvaluatorWithBaseImpl<Traits>("Dirichlet Boundary Conditions")
{
  m_dirichlets.clear();
  for( auto d: ds )
  {
    if( !d.isActive() ) continue;
    auto v = d.getValue();
    for( auto dof: d.m_local_dofs )
    {
      m_dirichlets.insert( std::make_pair(dof, v) );
    }
  }
};

template<typename Traits>
void DirichletsEvalutor<panzer::Traits::Hessian, Traits> :: preEvaluate(typename Traits::PreEvalData d)
{
  m_GhostedContainer = Teuchos::rcp_dynamic_cast<panzer::LinearObjContainer>(d.gedc->getDataObject("Ghosted Container"));
}

template<typename Traits>
void DirichletsEvalutor<panzer::Traits::Hessian, Traits> :: evaluateFields(const Teuchos::RCP<panzer::LinearObjContainer> thGhostedContainer)
{}
#endif

}

#endif
